#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

"""
@File    : visual_map.py
@Author  : Link
@Time    : 2022/6/5 10:31
@Mark    : 
"""
import math

import numpy as np
import pandas as pd

from PySide2 import QtWidgets, QtCore

from data_core.li import Li

from chart_core.chart_pyqtgraph.utils import GraphicsLayoutWidget, ImageItem, colormap, ColorBarItem, LabelItem
from chart_core.chart_pyqtgraph.core.visual import *


class VisualMapChart(QtWidgets.QWidget):
    key = None
    li = None

    def __init__(self):
        super(VisualMapChart, self).__init__()
        self.pw = GraphicsLayoutWidget()
        self.gridLayout = QtWidgets.QVBoxLayout(self)
        self.label = QtWidgets.QLabel()
        font = self.label.font()
        font.setPointSize(15)
        self.label.setFont(font)
        self.gridLayout.addWidget(self.label)
        self.gridLayout.addWidget(self.pw)
        # self.setFixedSize(900, 800)

    def set_li(self, li: Li):
        self.li = li
        self.li.chart_df_signal.connect(self.set_chart_df_chart)

    def set_data(self, key: str):
        if self.li is None:
            raise Exception("first set li")
        self.key = key
        self.set_title(key)

    def widget(self):
        return self

    def pw_show(self):
        self.show()

    def set_title(self, title: str = "PlotItem with CustomAxisItem, CustomViewBox"):
        self.setWindowTitle(title)

    def init_coord(self):
        x, y = self.li.chart_df["X_COORD"].min(), self.li.chart_df["Y_COORD"].min()
        if x == 0 and y == 0:
            return
        x1, y1 = self.li.df["X_COORD"].min(), self.li.df["Y_COORD"].min()
        if x < 0:
            self.li.chart_df["X_COORD"] = self.li.chart_df["X_COORD"] + abs(x1)
        if y < 0:
            self.li.chart_df["Y_COORD"] = self.li.chart_df["Y_COORD"] + abs(y1)
        if x > 0:
            self.li.chart_df["X_COORD"] = self.li.chart_df["X_COORD"] - x1
        if y > 0:
            self.li.chart_df["Y_COORD"] = self.li.chart_df["Y_COORD"] - y1

    def set_chart_df_chart(self):
        if self.key not in self.li.chart_df:
            return
        self.label.setText(self.key)
        x, y = self.li.chart_df["X_COORD"].max() + 1, self.li.chart_df["Y_COORD"].max() + 1
        group = self.li.chart_df.groupby("GROUP")
        self.pw.clear()
        if len(group) > 25:
            return
        row = math.ceil(math.sqrt(len(group)))
        items = []
        for index, each in enumerate(group):
            key, e_df = each  # type:str, pd.DataFrame
            data = np.full([x, y], np.nan)
            # e_df.apply(VisualMapChart.pd_coord_to_np_array, axis=1, args=(self.key, data))
            coord_to_np(e_df.X_COORD.to_numpy(), e_df.Y_COORD.to_numpy(), e_df[self.key].to_numpy(), data)
            t_row, t_col = divmod(index, row)
            im = ImageItem(image=data)
            items.append(im)
            p = self.pw.addPlot(t_row, t_col, 1, 1, title=key)
            p.getViewBox().invertY(True)
            p.addItem(im)
            p.setMouseEnabled(x=False, y=False)
        cmap = colormap.get('./colors/CET-D8')
        _min, _max = self.li.df[self.key].min(), self.li.df[self.key].max()
        rounding = (_max - _min) / 1E9
        bar = ColorBarItem(
            values=(self.li.df[self.key].quantile(0.05), self.li.df[self.key].quantile(0.95)),
            limits=(_min, _max),  # start with full range...
            rounding=rounding,
            cmap=cmap)
        bar.setImageItem(items)
        # manually adjust reserved space at top and bottom to align with plot
        bar.getAxis('bottom').setHeight(21)
        bar.getAxis('top').setHeight(31)
        self.pw.addItem(bar, 0, row + 1, 2, 1)  # large bar spanning both rows

    @staticmethod
    def pd_coord_to_np_array(row, key, data):
        data[row["X_COORD"]][row["Y_COORD"]] = row[key]

    def close_s(self):
        try:
            self.li.chart_df_signal.disconnect(self.set_chart_df_chart)
        except:
            pass
        self.li = None
